import copy
import random

import numpy as np


# 并查集实现
class Union_Set(object):

    def __init__(self, arr):
        self.parent = {pos: pos for pos in arr}
        self.count = len(arr)

    # 查找
    def find(self, root):
        if root == self.parent[root]:
            return root
        return self.find(self.parent[root])

    # 合并
    def union(self, root1, root2):
        self.parent[self.find(root1)] = self.find(root2)


# 迷宫生成类
class Maze_(object):

    def __init__(self, width=11, height=11):
        assert width >= 5 and height >= 5, \
            "Length of width or height must be larger than 5."

        self.width = (width // 2) * 2 + 1
        self.height = (height // 2) * 2 + 1
        self.start = [1, 0]
        self.destination = [self.height - 2, self.width - 1]
        self.matrix = None
        self.path = []

    # 打印矩阵
    def print_matrix(self):
        matrix = copy.deepcopy(self.matrix)
        for p in self.path:
            matrix[p[0]][p[1]] = 1
        for i in range(self.height):
            for j in range(self.width):
                if matrix[i][j] == -1:
                    print('□', end='')
                elif matrix[i][j] == 0:
                    print('  ', end='')
                elif matrix[i][j] == 1:
                    print('■', end='')
                elif matrix[i][j] == 2:
                    print('▲', end='')
            print('')

    # 生成矩阵
    def generate_matrix(self, mode, new_matrix):
        assert mode in [-1, 0, 1, 2, 3], \
            "Mode {} does not exist.".format(mode)

        if mode == -1:
            self.matrix = new_matrix
        elif mode == 0:
            self.generate_matrix_kruskal()
        elif mode == 1:
            self.generate_matrix_dfs()
        elif mode == 2:
            self.generate_matrix_prim()
        elif mode == 3:
            self.generate_matrix_split()

    # 重置矩阵大小
    def resize_matrix(self, width, height, mode, new_matrix):
        self.path = []
        self.width = (width // 2) * 2 + 1
        self.height = (height // 2) * 2 + 1
        self.start = [1, 0]
        self.destination = [self.height - 2, self.width - 1]
        self.generate_matrix(mode, new_matrix)

    # dfs算法生成矩阵
    def generate_matrix_dfs(self):
        # 地图初始化，并将出口和入口处的值设置为0
        self.matrix = -np.ones((self.height, self.width))
        self.matrix[self.start[0], self.start[1]] = 0
        self.matrix[self.destination[0], self.destination[1]] = 0

        visit_flag = [[0 for _ in range(self.width)] for _ in range(self.height)]

        # 检查
        def check(row1, col1):
            temp_sum = 0
            for lt in [[0, 1], [0, -1], [1, 0], [-1, 0]]:
                temp_sum += self.matrix[row1 + lt[0]][col1 + lt[1]]
            return temp_sum <= -3

        # dfs算法
        def dfs(row, col):
            visit_flag[row][col] = 1
            self.matrix[row][col] = 0
            if row == self.start[0] and col == self.start[1] + 1:
                return

            directions = [[0, 2], [0, -2], [2, 0], [-2, 0]]
            random.shuffle(directions)
            for lt in directions:
                row1, col1 = row + lt[0], col + lt[1]
                if 0 < row1 < self.height - 1 and 0 < col1 < self.width - 1 \
                        and visit_flag[row1][col1] == 0 and check(row1, col1):
                    if row == row1:
                        visit_flag[row][min(col, col1) + 1] = 1
                        self.matrix[row][min(col, col1) + 1] = 0
                    else:
                        visit_flag[min(row, row1) + 1][col] = 1
                        self.matrix[min(row, row1) + 1][col] = 0
                    dfs(row1, col1)

        dfs(self.destination[0], self.destination[1] - 1)
        self.matrix[self.start[0], self.start[1] + 1] = 0

    # primer算法生成矩阵
    def generate_matrix_prim(self):
        # 地图初始化，并将出口和入口处的值设置为0
        self.matrix = -np.ones((self.height, self.width))

        # 检查
        def check(rows, cols):
            temp_sum = 0
            for lt1 in [[0, 1], [0, -1], [1, 0], [-1, 0]]:
                temp_sum += self.matrix[rows + lt1[0]][cols + lt1[1]]
            return temp_sum < -3

        queue = []  # queue:队列
        row = (np.random.randint(1, self.height - 1) // 2) * 2 + 1
        col = (np.random.randint(1, self.width - 1) // 2) * 2 + 1
        queue.append((row, col, -1, -1))
        while len(queue) != 0:
            row, col, p, q = queue.pop(np.random.randint(0, len(queue)))
            if check(row, col):
                self.matrix[row, col] = 0
                if p != -1 and row == p:
                    self.matrix[row][min(col, q) + 1] = 0
                elif p != -1 and col == q:
                    self.matrix[min(row, p) + 1][col] = 0
                for lt in [[0, 2], [0, -2], [2, 0], [-2, 0]]:
                    row1, col1 = row + lt[0], col + lt[1]
                    if 0 < row1 < self.height - 1 and 0 < col1 < self.width - 1 \
                            and self.matrix[row1][col1] == -1:
                        queue.append((row1, col1, row, col))

        self.matrix[self.start[0], self.start[1]] = 0
        self.matrix[self.destination[0], self.destination[1]] = 0

    # 分割法生成矩阵
    def generate_matrix_split(self):
        # 地图初始化，并将出口和入口处的值设置为0
        self.matrix = -np.zeros((self.height, self.width))
        self.matrix[0, :] = -1
        self.matrix[self.height - 1, :] = -1
        self.matrix[:, 0] = -1
        self.matrix[:, self.width - 1] = -1

        # 随机生成偶数
        def get_random(start, end):
            rand = np.random.randint(start, end)
            if rand & 0x1 == 0:
                return rand
            return get_random(start, end)

        # split(左上角行数,列数,右下角行数,列数),墙壁在偶数行、列
        def split(left_row, left_col, right_row, right_col):
            if right_row - left_row < 2 or right_col - left_col < 2:
                return

            # 生成墙壁(偶数点)
            cur_row = get_random(left_row, right_row)
            cur_col = get_random(left_col, right_col)
            for i in range(left_col, right_col + 1):
                self.matrix[cur_row][i] = -1
            for i in range(left_row, right_row + 1):
                self.matrix[i][cur_col] = -1

            # 挖墙生成连通图(偶数点)
            wall_list = [
                ("left", cur_row, [left_col + 1, cur_col - 1]),
                ("right", cur_row, [cur_col + 1, right_col - 1]),
                ("top", cur_col, [left_row + 1, cur_row - 1]),
                ("down", cur_col, [cur_row + 1, right_row - 1])]

            random.shuffle(wall_list)
            for wall in wall_list[:-1]:
                if wall[2][1] - wall[2][0] < 1:
                    continue
                if wall[0] in ["left", "right"]:
                    self.matrix[wall[1], get_random(wall[2][0], wall[2][1] + 1) + 1] = 0
                else:
                    self.matrix[get_random(wall[2][0], wall[2][1] + 1), wall[1] + 1] = 0

            # 递归
            split(left_row + 2, left_col + 2, cur_row - 2, cur_col - 2)
            split(left_row + 2, cur_col + 2, cur_row - 2, right_col - 2)
            split(cur_row + 2, left_col + 2, right_row - 2, cur_col - 2)
            split(cur_row + 2, cur_col + 2, right_row - 2, right_col - 2)

            self.matrix[self.start[0], self.start[1]] = 0
            self.matrix[self.destination[0], self.destination[1]] = 0

        split(0, 0, self.height - 1, self.width - 1)

    # 最小生成树算法-kruskal(选边法)生成矩阵
    def generate_matrix_kruskal(self):
        # 地图初始化，并将出口和入口处的值设置为0
        self.matrix = -np.ones((self.height, self.width))

        # 检查
        def check(rows, cols):
            ans, counter = [], 0
            for type_list in [[0, 1], [0, -1], [1, 0], [-1, 0]]:
                row1_, col1_ = rows + type_list[0], cols + type_list[1]
                if 0 < row1_ < self.height - 1 and 0 < col1_ < self.width - 1 \
                        and self.matrix[row1_, col1_] == -1:
                    ans.append([type_list[0] * 2, type_list[1] * 2])
                    counter += 1
            if counter <= 1:
                return []
            return ans

        nodes = set()
        row = 1
        while row < self.height:
            col = 1
            while col < self.width:
                self.matrix[row, col] = 0
                nodes.add((row, col))
                col += 2
            row += 2

        union_set = Union_Set(nodes)
        while union_set.count > 1:
            row, col = nodes.pop()
            directions = check(row, col)
            if len(directions):
                random.shuffle(directions)
                for lt in directions:
                    row1, col1 = row + lt[0], col + lt[1]
                    if union_set.find((row, col)) == union_set.find((row1, col1)):
                        continue
                    nodes.add((row, col))
                    union_set.count -= 1
                    union_set.union((row, col), (row1, col1))

                    if row == row1:
                        self.matrix[row][min(col, col1) + 1] = 0
                    else:
                        self.matrix[min(row, row1) + 1][col] = 0
                    break

        self.matrix[self.start[0], self.start[1]] = 0
        self.matrix[self.destination[0], self.destination[1]] = 0

    # 迷宫寻路dfs算法
    def find_path_dfs(self, destination):
        visited = [[0 for _ in range(self.width)] for _ in range(self.height)]
        global count1
        count1 = 0

        # 定义dfs寻路算法
        def dfs(path):
            global count1
            visited[path[-1][0]][path[-1][1]] = 1
            if path[-1][0] == destination[0] and path[-1][1] == destination[1]:
                self.path = path[:]
                return count1
            for lt in [[0, 1], [0, -1], [1, 0], [-1, 0]]:
                row1, col1 = path[-1][0] + lt[0], path[-1][1] + lt[1]
                if 0 < row1 < self.height - 1 and 0 < col1 < self.width \
                        and visited[row1][col1] == 0 and self.matrix[row1][col1] == 0:
                    count1 += 1
                    dfs(path + [[row1, col1]])

        dfs([[self.start[0], self.start[1]]])


# 迷宫生成部分主函数
if __name__ == '__main__':
    maze = Maze_(51, 51)
    count1 = 0
    maze.generate_matrix_prim()
    maze.print_matrix()
    maze.find_path_dfs(maze.destination)
    print("answer", maze.path)
    maze.print_matrix()
